/*
    Simple Madryga encryption algorithm
    Author: @5mukx
*/

use std::fs::{self, File};
use std::io::{self, Read, Write};
use std::path::{Path};

const BLOCK_SIZE: usize = 8;
const NUM_ROUNDS: usize = 8;

const RANDOM_CONSTANT: u64 = 0x0f1e2d3c4b5a6978;

fn rotate_left_16(value: u16, bits: u32) -> u16 {
    // left 
    (value << bits) | (value >> ((16 - bits) % 16))
}

fn rotate_right_16(value: u16, bits: u32) -> u16 {
    (value >> bits) | (value << ((16 - bits) % 16))
}

fn madryga_inner_cycle_encrypt(data: &mut [u8], key: u64){
    let data_len = data.len();
    
    for i in(0..data_len - 1).rev(){
        let start = i % data_len;
        let mid = (i + 1) % data_len;
        let end = (i + 2) % data_len;

        let rotation_bits = (data[end] & 0x07) as u32;

        let mut temp: u16 = ((data[mid] as u16) << 8) | (data[start] as u16);

        temp = rotate_left_16(temp, rotation_bits);
        data[start] = temp as u8;
        data[mid] = (temp >> 8) as u8;

        data[end] ^= (key & 0xFF) as u8;
    }
}

fn madryga_inner_cycle_decrypt(data: &mut [u8], key: u64) {
    let data_len = data.len();
    for i in 0..data_len - 1 {
        let start = i % data_len;
        let mid = (i + 1) % data_len;
        let end = (i + 2) % data_len;

        data[end] ^= (key & 0xFF) as u8;

        let rotation_bits = (data[end] & 0x07) as u32;

        let mut temp: u16 = ((data[mid] as u16) << 8) | (data[start] as u16);
        temp = rotate_right_16(temp, rotation_bits);
        data[start] = temp as u8;
        data[mid] = (temp >> 8) as u8;
    }
}

fn madryga_encrypt(block: &mut [u8], mut key: u64) {
    key ^= RANDOM_CONSTANT;
    for _ in 0..NUM_ROUNDS {
        madryga_inner_cycle_encrypt(block, key);
    }
}

fn madryga_decrypt(block: &mut [u8], mut key: u64) {
    key ^= RANDOM_CONSTANT;
    for _ in 0..NUM_ROUNDS {
        madryga_inner_cycle_decrypt(block, key);
    }
}

fn encrypt_file(input_path: &str, output_path: &str, key: u64) -> io::Result<()> {
    let mut input_file = File::open(input_path)?;
    let mut output_file = File::create(output_path)?;

    let mut buffer = [0u8; BLOCK_SIZE];
    loop {
        let bytes_read = input_file.read(&mut buffer)?;
        if bytes_read == 0 {
            break;
        }
        if bytes_read < BLOCK_SIZE {
            for i in bytes_read..BLOCK_SIZE {
                buffer[i] = 0x90;
            }
        }
        madryga_encrypt(&mut buffer, key);
        output_file.write_all(&buffer)?;
    }
    Ok(())
}

fn decrypt_file(input_path: &str, output_path: &str, key: u64) -> io::Result<()> {
    let mut input_file = File::open(input_path)?;
    let mut output_file = File::create(output_path)?;

    let metadata = input_file.metadata()?;
    let file_size = metadata.len() as usize;
    if file_size % BLOCK_SIZE != 0 {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "File size is not a multiple of block size",
        ));
    }

    let mut buffer = [0u8; BLOCK_SIZE];
    let mut remaining = file_size;
    while remaining > BLOCK_SIZE {
        input_file.read_exact(&mut buffer)?;
        madryga_decrypt(&mut buffer, key);
        output_file.write_all(&buffer)?;
        remaining -= BLOCK_SIZE;
    }

    if remaining > 0 {
        input_file.read_exact(&mut buffer)?;
        madryga_decrypt(&mut buffer, key);
        let mut len = BLOCK_SIZE;
        while len > 0 && buffer[len - 1] == 0x90 {
            len -= 1;
        }
        output_file.write_all(&buffer[0..len])?;
    }

    Ok(())
}

fn handle_files(folder_path:&Path, key:u64) -> io::Result<()>{
    
    for entry in fs::read_dir(folder_path)? {
        let entry = entry?;
        let path = entry.path();
        let file_name_os = path.file_name().unwrap_or_default();
        let file_name = file_name_os.to_str().unwrap_or_default();
        if file_name == "." || file_name == ".." {
            continue;
        }
        if path.is_dir() {
            handle_files(&path, key)?;
        } else {
            println!("File: {}", path.display());
            let new_file_name = format!("{}.bin", file_name);
            let encrypted_path = path.with_file_name(new_file_name);
            encrypt_file(
                path.to_str().unwrap(),
                encrypted_path.to_str().unwrap(),
                key,
            )?;
        }
    }
    
    Ok(())
}

fn decrypt_files(folder_path: &Path, key: u64) -> io::Result<()> {
    for entry in fs::read_dir(folder_path)? {
        let entry = entry?;
        let path = entry.path();
        let file_name_os = path.file_name().unwrap_or_default();
        let file_name = file_name_os.to_str().unwrap_or_default();
        if file_name == "." || file_name == ".." {
            continue;
        }
        if path.is_dir() {
            decrypt_files(&path, key)?;
        } else if file_name.ends_with(".bin") {
            println!("File: {}", path.display());
            let new_file_name = format!("{}.decrypted", file_name);
            let decrypted_path = path.with_file_name(new_file_name);
            decrypt_file(
                path.to_str().unwrap(),
                decrypted_path.to_str().unwrap(),
                key,
            )?;
        }
    }
    Ok(())
}


fn main()-> io::Result<()> {
    let root_folder = Path::new(r".\test");
    let key: u64 = 0x0123456789ABCDEF;

    handle_files(root_folder, key)?;
    decrypt_files(root_folder, key)?;

    Ok(())
}


